import asyncio
import websockets
import json
import logging

class WebSocketClient:
    def __init__(self, uri, access_token, device_id, protocol_version=3):
        self.uri = uri
        self.access_token = access_token
        self.device_id = device_id
        self.protocol_version = protocol_version
        self.websocket = None
    

    async def connect(self):
        """连接到 WebSocket 服务器"""
        headers = {
            "Authorization": f"Bearer {self.access_token}",
            "Device-Id": self.device_id,
            "Protocol-Version": str(self.protocol_version)
        }
        
        logging.info(f"Connecting to WebSocket server at {self.uri}...")
        self.websocket = await websockets.connect(self.uri, extra_headers=headers)
        logging.info("Connected to WebSocket server.")

    async def send_message(self, message):
        """发送消息到服务器"""
        if self.websocket and self.websocket.open:
            await self.websocket.send(message)
            logging.info(f">>> Sent message to server: {message}")
        else:
            logging.info("WebSocket is not connected.")

    async def send_status(self,status_str:str):
        message = {
            "type": "state",
            "state": status_str
        }
        json_message = json.dumps(message)
        self.ws.send(json_message)
        logging.info(f"Sent status: {json_message}")

    async def receive_message(self):
        """接收来自服务器的消息并解析"""
        if self.websocket and self.websocket.open:
            response = await self.websocket.recv()
            print(f"<<< Received message from server: {response}")  # 打印接收到的消息
            try:
                parsed_message = json.loads(response)
                return parsed_message
            except json.JSONDecodeError as e:
                print(f"Failed to parse JSON message: {e}")
                return None
        else:
            print("WebSocket is not connected.")
            return None

    async def close(self):
        """关闭连接"""
        if self.websocket and self.websocket.open:
            await self.websocket.close()
            print("Connection closed.")
        else:
            print("WebSocket is already closed or not connected.")

    async def run_forever(self):
        """持续监听服务器消息"""
        try:
            while True:
                message = await self.receive_message()
                if message is None:
                    break
        except KeyboardInterrupt:
            print("Interrupted. Closing connection...")
        except Exception as e:
            print(f"An error occurred: {e}")
        finally:
            await self.close()
